/**
 * Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
/*!
* \file cube_in_buffer_params.h
* \brief
*/

#ifndef IMPL_MATMUL_MODULES_RESOURCE_CUBE_IN_BUFFER_CUBE_IN_BUFFER_PARAMS_H_
#define IMPL_MATMUL_MODULES_RESOURCE_CUBE_IN_BUFFER_CUBE_IN_BUFFER_PARAMS_H_

#include "cube_in_buffer_utils.h"
#include "../../feature_trait/matmul_feature_trait.h"

namespace matmul {
constexpr int32_t BANK_CONFLICT_SIZE = 512;
template <typename IMPL, const auto& MM_CFG, class INPUT_TYPE, typename = void>
class CubeInBufferParams {
public:
    __aicore__ inline int32_t GetMajorCacheNum()
    {
        if constexpr (DoMatmulMDL(MM_CFG) || DoMatmulSpecialMDL(MM_CFG)) {
            return MATMUL_PARAM_VAR.tiling_.GetStepKa();
        } else if constexpr (DoMatmulSpecialBasicBlock(MM_CFG)) {
            return ToMatmulConfig(MM_CFG).stepM;
        } else {
            return MATMUL_PARAM_VAR.tiling_.GetStepM();
        }
    }

    __aicore__ inline int32_t GetTotalCacheNum()
    {
        if constexpr (DoMatmulMDL(MM_CFG) || DoMatmulSpecialMDL(MM_CFG)) {
            return MATMUL_PARAM_VAR.tiling_.GetStepKa() * MATMUL_PARAM_VAR.tiling_.GetStepM();
        } else {
            return 1;
        }
    }

    __aicore__ inline static int32_t GetCurKPos(int32_t curRow, int32_t curCol)
    {
        return curCol;
    }

    __aicore__ inline static int32_t GetCurMNPos(int32_t curRow, int32_t curCol)
    {
        return curRow;
    }

    __aicore__ inline int32_t GetInnerLoopCnt()
    {
        return MATMUL_PARAM_VAR.kIter_;
    }

    __aicore__ inline int32_t GetIterIndex(int32_t curRow, int32_t curCol)
    {
        if constexpr (DoMatmulNorm(MM_CFG) || DoMatmulIBShareNorm(MM_CFG) || DoMatmulBasicBlock(MM_CFG)) {
            auto& var = MATMUL_PARAM_VAR;
            if (var.tiling_.GetIterateOrder() == static_cast<int>(IterateOrder::ORDER_M)) {
                return curCol;
            } else {
                return (curRow * var.kIter_ + curCol) % (var.tiling_.GetStepM() * var.kIter_);
            }
        } else if constexpr (DoMatmulSpecialBasicBlock(MM_CFG)){
            auto& var = MATMUL_PARAM_VAR;
            if (var.tiling_.GetIterateOrder() == static_cast<int>(IterateOrder::ORDER_M)) {
                return curCol;
            } else {
                return (curRow * var.kIter_ + curCol) %
                    (ToMatmulConfig(MM_CFG).stepM * ToMatmulConfig(MM_CFG).singleCoreK / ToMatmulConfig(MM_CFG).basicK);
            }
        } else {
            return 0;
        }
    }

    __aicore__ inline bool IsTailBlock()
    {
        auto& var = MATMUL_PARAM_VAR;
        return (var.baseUseM_ != var.tiling_.GetBaseM()) || (var.baseUseK_ != var.tiling_.GetBaseK());
    }

    __aicore__ inline int32_t GetBankConflictSize()
    {
        if constexpr (MatmulFeatureTrait<MM_CFG>::IsNeedUB() &&
            ToMatmulConfig(MM_CFG).enVecND2NZ && (DoMatmulMDL(MM_CFG) || DoMatmulSpecialMDL(MM_CFG))) {
            if constexpr (INPUT_TYPE::format == CubeFormat::ND) {
                auto& var = MATMUL_PARAM_VAR;
                constexpr int32_t c0Size = AuxGetC0Size<typename INPUT_TYPE::TRANS_T>();
                if constexpr (INPUT_TYPE::isTrans) {
                    bool isBankConflict = Ceil(var.tiling_.GetStepM() * var.tiling_.GetBaseM(), c0Size) *
                        ONE_BLK_SIZE % BANK_CONFLICT_SIZE == 0;
                    return isBankConflict ? var.tiling_.GetBaseK() * c0Size *
                        var.tiling_.GetStepKa() * sizeof(typename INPUT_TYPE::TRANS_T) : 0;
                } else {
                    bool isBankConflict = Ceil(var.tiling_.GetStepKa() * var.tiling_.GetBaseK(), c0Size) *
                        ONE_BLK_SIZE % BANK_CONFLICT_SIZE == 0;
                    return isBankConflict ? var.tiling_.GetBaseM() * c0Size *
                        var.tiling_.GetStepM() * sizeof(typename INPUT_TYPE::TRANS_T) : 0;
                }
            } else {
                return 0;
            }
        } else {
            return 0;
        }
    }

    __aicore__ inline bool IsDataAmountEqual()
    {
        auto& var = MATMUL_PARAM_VAR;
        return (var.baseUseM_ == var.tiling_.GetBaseM()) &&
            (var.baseUseK_ == var.tiling_.GetBaseK());
    }
};

template <typename IMPL, const auto& MM_CFG, class INPUT_TYPE>
class CubeInBufferParams<IMPL, MM_CFG, INPUT_TYPE, enable_if_t<INPUT_TYPE::TAG == InputTypeTag::B>> {
public:
    __aicore__ inline int32_t GetMajorCacheNum()
    {
        if constexpr (DoMatmulMDL(MM_CFG) || DoMatmulSpecialMDL(MM_CFG)) {
            return MATMUL_PARAM_VAR.tiling_.GetStepKb();
        } else if constexpr (DoMatmulSpecialBasicBlock(MM_CFG)) {
            return ToMatmulConfig(MM_CFG).stepN;
        } else {
            return MATMUL_PARAM_VAR.tiling_.GetStepN();
        }
    }

    __aicore__ inline int32_t GetTotalCacheNum()
    {
        if constexpr (DoMatmulMDL(MM_CFG) || DoMatmulSpecialMDL(MM_CFG) ||
            (DoMatmulNorm(MM_CFG) && ToMatmulConfig(MM_CFG).intraBlockPartSum)) {
            return MATMUL_PARAM_VAR.tiling_.GetStepKb() * MATMUL_PARAM_VAR.tiling_.GetStepN();
        } else {
            return 1;
        }
    }

    __aicore__ inline static int32_t GetCurKPos(int32_t curRow, int32_t curCol)
    {
        return curRow;
    }

    __aicore__ inline static int32_t GetCurMNPos(int32_t curRow, int32_t curCol)
    {
        return curCol;
    }

    __aicore__ inline int32_t GetInnerLoopCnt()
    {
        return MATMUL_PARAM_VAR.kIter_;
    }

    __aicore__ inline int32_t GetIterIndex(int32_t curRow, int32_t curCol)
    {
        if constexpr (DoMatmulNorm(MM_CFG) || DoMatmulIBShareNorm(MM_CFG) || DoMatmulBasicBlock(MM_CFG)) {
            auto& var = MATMUL_PARAM_VAR;
            if (var.tiling_.GetIterateOrder() == static_cast<int>(IterateOrder::ORDER_M)) {
                return (curRow + curCol * var.kIter_) % (var.tiling_.GetStepN() * var.kIter_);
            } else {
                return curRow;
            }
        } else if (DoMatmulSpecialBasicBlock(MM_CFG)) {
             auto& var = MATMUL_PARAM_VAR;
            if (var.tiling_.GetIterateOrder() == static_cast<int>(IterateOrder::ORDER_M)) {
                return (curRow + curCol * var.kIter_) %
                    ToMatmulConfig(MM_CFG).stepN * ToMatmulConfig(MM_CFG).singleCoreK / ToMatmulConfig(MM_CFG).basicK;
            } else {
                return curRow;
            }
        } else {
            return 0;
        }
    }

    __aicore__ inline bool IsTailBlock()
    {
        auto& var = MATMUL_PARAM_VAR;
        return (var.baseUseN_ != var.tiling_.GetBaseN()) || (var.baseUseK_ != var.tiling_.GetBaseK());
    }

    __aicore__ inline int32_t GetBankConflictSize()
    {
        if constexpr (MatmulFeatureTrait<MM_CFG>::IsNeedUB() &&
            ToMatmulConfig(MM_CFG).enVecND2NZ && (DoMatmulMDL(MM_CFG) || DoMatmulSpecialMDL(MM_CFG))) {
            if constexpr (INPUT_TYPE::format == CubeFormat::ND) {
                auto& var = MATMUL_PARAM_VAR;
                constexpr int32_t c0Size = AuxGetC0Size<typename INPUT_TYPE::TRANS_T>();
                if constexpr (INPUT_TYPE::isTrans) {
                    bool isBankConflict = Ceil(var.tiling_.GetStepKb() * var.tiling_.GetBaseK(), c0Size) *
                        ONE_BLK_SIZE % BANK_CONFLICT_SIZE == 0;
                    return isBankConflict ? var.tiling_.GetBaseN() * c0Size *
                        var.tiling_.GetStepN() * sizeof(typename INPUT_TYPE::TRANS_T) : 0;
                } else {
                    bool isBankConflict = Ceil(var.tiling_.GetStepN() * var.tiling_.GetBaseN(), c0Size) *
                        ONE_BLK_SIZE % BANK_CONFLICT_SIZE == 0;
                    return isBankConflict ? var.tiling_.GetBaseK() * c0Size *
                        var.tiling_.GetStepKb() * sizeof(typename INPUT_TYPE::TRANS_T) : 0;
                }
            } else {
                return 0;
            }
        } else {
            return 0;
        }
    }

    __aicore__ inline bool IsDataAmountEqual()
    {
        auto& var = MATMUL_PARAM_VAR;
        return (var.baseUseK_ == var.tiling_.GetBaseK()) &&
            (var.baseUseN_ == var.tiling_.GetBaseN());
    }
};

}
#endif // _CUBE_IN_BUFFER_PARAMS_H_